"""
MongoDB工具模块
"""
import logging
import pandas as pd
from pymongo import MongoClient
from pymongo.errors import ConnectionFailure, OperationFailure

logger = logging.getLogger(__name__)

class MongoUtils:
    """MongoDB工具类"""
    
    @staticmethod
    def create_client_from_config(config):
        """
        根据配置创建MongoDB客户端
        
        Args:
            config: 数据源配置
        
        Returns:
            pymongo.MongoClient: MongoDB客户端
        """
        host = config.get('host')
        port = config.get('port')
        username = config.get('username')
        password = config.get('password')
        database = config.get('database')
        auth_source = config.get('auth_source', 'admin')
        
        connection_string = None
        
        # 构建连接字符串
        if username and password:
            connection_string = f"mongodb://{username}:{password}@{host}:{port}/{database}?authSource={auth_source}"
        else:
            connection_string = f"mongodb://{host}:{port}/{database}"
        
        try:
            client = MongoClient(connection_string)
            return client
        except Exception as e:
            logger.error(f"创建MongoDB客户端失败: {str(e)}")
            raise
    
    @staticmethod
    def test_connection(config):
        """
        测试MongoDB连接
        
        Args:
            config: 数据源配置
        
        Returns:
            bool: 连接是否成功
        """
        try:
            client = MongoUtils.create_client_from_config(config)
            # 执行简单命令以验证连接
            client.admin.command('ping')
            return True
        except (ConnectionFailure, OperationFailure) as e:
            logger.error(f"测试MongoDB连接失败: {str(e)}")
            return False
        finally:
            if 'client' in locals():
                client.close()
    
    @staticmethod
    def get_collections(client, database_name):
        """
        获取所有集合名称
        
        Args:
            client: MongoDB客户端
            database_name: 数据库名称
        
        Returns:
            list: 集合名称列表
        """
        try:
            db = client[database_name]
            return db.list_collection_names()
        except Exception as e:
            logger.error(f"获取集合列表失败: {str(e)}")
            raise
    
    @staticmethod
    def get_collection_schema(client, database_name, collection_name, sample_size=100):
        """
        获取集合结构（基于采样）
        
        Args:
            client: MongoDB客户端
            database_name: 数据库名称
            collection_name: 集合名称
            sample_size: 采样大小
        
        Returns:
            dict: 字段信息字典
        """
        try:
            db = client[database_name]
            collection = db[collection_name]
            
            # 获取样本文档
            sample_docs = list(collection.find().limit(sample_size))
            
            if not sample_docs:
                return {}
            
            # 分析字段类型
            schema = {}
            for doc in sample_docs:
                for field, value in doc.items():
                    if field not in schema:
                        schema[field] = {'type': type(value).__name__, 'sample': value}
            
            return schema
        except Exception as e:
            logger.error(f"获取集合结构失败: {str(e)}")
            raise
    
    @staticmethod
    def find_documents(client, database_name, collection_name, query=None, projection=None, limit=0, skip=0, sort=None):
        """
        查询文档
        
        Args:
            client: MongoDB客户端
            database_name: 数据库名称
            collection_name: 集合名称
            query: 查询条件 (可选)
            projection: 投影 (可选)
            limit: 限制数量 (可选)
            skip: 跳过数量 (可选)
            sort: 排序 (可选)
        
        Returns:
            list: 文档列表
        """
        try:
            db = client[database_name]
            collection = db[collection_name]
            
            cursor = collection.find(
                filter=query or {},
                projection=projection,
            )
            
            if skip:
                cursor = cursor.skip(skip)
            
            if limit:
                cursor = cursor.limit(limit)
            
            if sort:
                cursor = cursor.sort(sort)
            
            return list(cursor)
        except Exception as e:
            logger.error(f"查询文档失败: {str(e)}")
            raise
    
    @staticmethod
    def find_to_df(client, database_name, collection_name, query=None, projection=None, limit=0, skip=0, sort=None):
        """
        查询文档并返回DataFrame
        
        Args:
            client: MongoDB客户端
            database_name: 数据库名称
            collection_name: 集合名称
            query: 查询条件 (可选)
            projection: 投影 (可选)
            limit: 限制数量 (可选)
            skip: 跳过数量 (可选)
            sort: 排序 (可选)
        
        Returns:
            pandas.DataFrame: 查询结果
        """
        try:
            documents = MongoUtils.find_documents(
                client, database_name, collection_name, 
                query, projection, limit, skip, sort
            )
            
            if not documents:
                return pd.DataFrame()
            
            return pd.DataFrame(documents)
        except Exception as e:
            logger.error(f"查询文档并转换为DataFrame失败: {str(e)}")
            raise
    
    @staticmethod
    def insert_documents(client, database_name, collection_name, documents):
        """
        插入文档
        
        Args:
            client: MongoDB客户端
            database_name: 数据库名称
            collection_name: 集合名称
            documents: 文档列表或单个文档
        
        Returns:
            int: 插入的文档数量
        """
        try:
            db = client[database_name]
            collection = db[collection_name]
            
            if isinstance(documents, list):
                if not documents:
                    return 0
                result = collection.insert_many(documents)
                return len(result.inserted_ids)
            else:
                result = collection.insert_one(documents)
                return 1 if result.inserted_id else 0
        except Exception as e:
            logger.error(f"插入文档失败: {str(e)}")
            raise
    
    @staticmethod
    def update_documents(client, database_name, collection_name, query, update, upsert=False, multi=True):
        """
        更新文档
        
        Args:
            client: MongoDB客户端
            database_name: 数据库名称
            collection_name: 集合名称
            query: 查询条件
            update: 更新内容
            upsert: 不存在时是否插入
            multi: 是否更新多个文档
        
        Returns:
            int: 更新的文档数量
        """
        try:
            db = client[database_name]
            collection = db[collection_name]
            
            if multi:
                result = collection.update_many(query, update, upsert=upsert)
                return result.modified_count
            else:
                result = collection.update_one(query, update, upsert=upsert)
                return result.modified_count
        except Exception as e:
            logger.error(f"更新文档失败: {str(e)}")
            raise
    
    @staticmethod
    def delete_documents(client, database_name, collection_name, query, multi=True):
        """
        删除文档
        
        Args:
            client: MongoDB客户端
            database_name: 数据库名称
            collection_name: 集合名称
            query: 查询条件
            multi: 是否删除多个文档
        
        Returns:
            int: 删除的文档数量
        """
        try:
            db = client[database_name]
            collection = db[collection_name]
            
            if multi:
                result = collection.delete_many(query)
                return result.deleted_count
            else:
                result = collection.delete_one(query)
                return result.deleted_count
        except Exception as e:
            logger.error(f"删除文档失败: {str(e)}")
            raise
    
    @staticmethod
    def count_documents(client, database_name, collection_name, query=None):
        """
        统计文档数量
        
        Args:
            client: MongoDB客户端
            database_name: 数据库名称
            collection_name: 集合名称
            query: 查询条件 (可选)
        
        Returns:
            int: 文档数量
        """
        try:
            db = client[database_name]
            collection = db[collection_name]
            
            return collection.count_documents(query or {})
        except Exception as e:
            logger.error(f"统计文档数量失败: {str(e)}")
            raise
    
    @staticmethod
    def aggregate(client, database_name, collection_name, pipeline):
        """
        聚合查询
        
        Args:
            client: MongoDB客户端
            database_name: 数据库名称
            collection_name: 集合名称
            pipeline: 聚合管道
        
        Returns:
            list: 聚合结果
        """
        try:
            db = client[database_name]
            collection = db[collection_name]
            
            return list(collection.aggregate(pipeline))
        except Exception as e:
            logger.error(f"聚合查询失败: {str(e)}")
            raise
    
    @staticmethod
    def df_to_mongo(df, client, database_name, collection_name, if_exists='replace'):
        """
        将DataFrame写入MongoDB
        
        Args:
            df: pandas.DataFrame
            client: MongoDB客户端
            database_name: 数据库名称
            collection_name: 集合名称
            if_exists: 如果集合存在的处理方式 ('fail', 'replace', 'append')
        
        Returns:
            int: 插入的文档数量
        """
        try:
            db = client[database_name]
            collection = db[collection_name]
            
            if if_exists == 'replace' and collection_name in db.list_collection_names():
                collection.drop()
            elif if_exists == 'fail' and collection_name in db.list_collection_names():
                raise ValueError(f"集合 {collection_name} 已存在")
            
            # 转换DataFrame为文档列表
            records = df.to_dict('records')
            
            if not records:
                return 0
            
            result = collection.insert_many(records)
            return len(result.inserted_ids)
        except Exception as e:
            logger.error(f"将DataFrame写入MongoDB失败: {str(e)}")
            raise